use serde::{Deserialize, Serialize};
use serde_with::serde_as;

use crate::{
    utils::const_assert::{ConstAssert, IsTrue},
    voxel::layout::{chunk_dim, chunk_voxel_count},
};

///
/// tile在chunk内的局部坐标
///

#[serde_as]
#[derive(Clone, Serialize, Deserialize)]
pub struct Chunk<const CBITS: u8, DATA: Clone + Serialize + serde::de::DeserializeOwned + Default>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    ///
    /// 每个顶点的密度信息，顶点坐标范围 [0, C)，其中 C = chunk_dim
    /// +X，+Y, +Z三个方向的halo边界数据由外部获取
    ///
    #[serde_as(as = "Box<[_; chunk_voxel_count::<CBITS>()]>")]
    pub densities: Box<[f32; chunk_voxel_count::<CBITS>()]>,

    pub materials: TileMaterial<CBITS>,

    #[serde(bound(
        serialize = "DATA: Serialize",
        deserialize = "DATA: serde::de::DeserializeOwned"
    ))]
    pub data: rustc_hash::FxHashMap<u32, DATA>,
}

impl<const CBITS: u8, DATA: Clone + Serialize + serde::de::DeserializeOwned + Default>
    Chunk<CBITS, DATA>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    const C: usize = chunk_dim::<CBITS>();

    pub fn new() -> Self {
        let mut densities_vec = Vec::with_capacity(chunk_voxel_count::<CBITS>());
        densities_vec.resize(chunk_voxel_count::<CBITS>(), 0.0f32);
        let densities: Box<[f32; chunk_voxel_count::<CBITS>()]> = match densities_vec.try_into() {
            Ok(b) => b,
            Err(_) => unreachable!(),
        };
        Self {
            densities,
            materials: TileMaterial::new_palette(),
            data: rustc_hash::FxHashMap::default(),
        }
    }

    #[inline]
    pub fn chunk_vertex_linear_index(ix: u32, iy: u32, iz: u32) -> usize {
        let (x, y, z) = (ix as usize, iy as usize, iz as usize);
        x + Self::C * (y + Self::C * z)
    }

    ///
    /// 设置密度，设置的密度永远是tile原点所在的顶点的密度（也是就tile(0,0,0)所在的顶点的密度）
    /// 顶点坐标范围: [0, C)，其中 C = chunk_dim
    ///
    #[inline]
    pub fn set_density(&mut self, ix: u32, iy: u32, iz: u32, density: f32) {
        // 边界检查：顶点坐标范围是 [0, C)
        let c = Self::C as u32;
        debug_assert!(ix < c && iy < c && iz < c, "set_density OOB");
        let index = Self::chunk_vertex_linear_index(ix, iy, iz);
        self.densities[index] = density;
    }

    #[inline]
    pub fn get_density(&self, ix: u32, iy: u32, iz: u32) -> f32 {
        let c = Self::C as u32;
        debug_assert!(ix < c && iy < c && iz < c, "get_density OOB");
        let index = Self::chunk_vertex_linear_index(ix, iy, iz);
        self.densities[index]
    }
}

#[serde_as]
#[derive(Clone, Serialize, Deserialize)]
pub struct ChunkPalette<const CBITS: u8>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    #[serde_as(as = "[_; 256]")]
    entries: [TileMaterialRef; 256],
    #[serde_as(as = "[_; 256]")]
    refcnt: [u32; 256],
    #[serde_as(as = "[_; chunk_voxel_count::<CBITS>()]")]
    index: [u8; chunk_voxel_count::<CBITS>()], // 索引到entries的索引
}

impl<const CBITS: u8> ChunkPalette<CBITS>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    pub fn new() -> Self {
        Self {
            entries: [TileMaterialRef::AIR; 256], // 使用默认材质，实际不会被使用
            refcnt: [0; 256],
            index: [255; chunk_voxel_count::<CBITS>()], // 255 表示空气/空
        }
    }

    #[inline]
    pub fn get_tile_material(&self, index: usize) -> Option<TileMaterialRef> {
        let id = *self.index.get(index)?;
        if id == 255 {
            return None;
        }
        let idx = id as usize;
        debug_assert!(self.refcnt[idx] > 0);
        Some(self.entries[idx])
    }

    fn find_slot(&self, m: TileMaterialRef) -> Option<u8> {
        // 只在 refcnt>0 的槽位里找
        for s in 0..255 {
            if self.refcnt[s] > 0 && self.entries[s] == m {
                return Some(s as u8);
            }
        }
        None
    }

    fn find_free_slot(&self) -> Option<u8> {
        for s in 0..255 {
            if self.refcnt[s] == 0 {
                return Some(s as u8);
            }
        }
        None
    }

    pub fn set_tile_material(&mut self, i: usize, m: Option<TileMaterialRef>) -> Result<(), ()> {
        let old = *self.index.get(i).ok_or(())?;
        // 先释放旧引用
        if old != 255 {
            let s = old as usize;
            self.refcnt[s] = self.refcnt[s].saturating_sub(1);
            // refcnt==0 时槽位可复用；entries[s] 保留原值以便未来命中（可选是否清零）
        }

        match m {
            None => {
                // 写空（空气）
                self.index[i] = 255;
                Ok(())
            }
            Some(mat) => {
                // 先看是否已有同材质的槽位
                if let Some(s) = self.find_slot(mat) {
                    self.index[i] = s;
                    self.refcnt[s as usize] += 1;
                    return Ok(());
                }
                // 找空槽
                if let Some(s) = self.find_free_slot() {
                    let si = s as usize;
                    self.entries[si] = mat;
                    self.refcnt[si] = 1;
                    self.index[i] = s;
                    Ok(())
                } else {
                    Err(()) // palette 满；由上层决定是否回退到 Direct
                }
            }
        }
    }

    pub fn to_direct(&self) -> TileMaterial<CBITS> {
        let mut v = Vec::with_capacity(chunk_voxel_count::<CBITS>());
        for &idx in &self.index {
            if idx == 255 {
                v.push(None);
            } else {
                v.push(Some(self.entries[idx as usize]));
            }
        }
        TileMaterial::Direct(v)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TileMaterialRef {
    ///
    /// 空材质
    ///
    AIR,
    ///
    /// 材质索引
    ///
    MATERIAL(u32),
}

#[derive(Clone, Serialize, Deserialize)]
pub enum TileMaterial<const CBITS: u8>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    ///
    /// 调色板，当一个Chunk中材质的种类<255时，使用调色板
    ///
    ChunkPalette(ChunkPalette<CBITS>),
    ///
    /// 直接使用材质索引，当一个Chunk中材质的种类>=255时，使用直接使用材质索引
    ///
    Direct(Vec<Option<TileMaterialRef>>),
}

impl<const CBITS: u8> TileMaterial<CBITS>
where
    [(); chunk_voxel_count::<CBITS>()]:,
    ConstAssert<{ CBITS <= 8 }>: IsTrue,
{
    pub fn new_palette() -> Self {
        TileMaterial::ChunkPalette(ChunkPalette::new())
    }

    #[inline]
    pub fn get_tile_material(&self, index: usize) -> Option<TileMaterialRef> {
        match self {
            TileMaterial::ChunkPalette(palette) => palette.get_tile_material(index),
            TileMaterial::Direct(direct) => {
                if let Some(mat) = direct.get(index).copied() {
                    mat
                } else {
                    None
                }
            }
        }
    }

    pub fn set_tile_material(&mut self, i: usize, m: Option<TileMaterialRef>) {
        match self {
            TileMaterial::ChunkPalette(p) => {
                if p.set_tile_material(i, m).is_err() {
                    let mut d = Vec::with_capacity(chunk_voxel_count::<CBITS>());
                    for t in 0..chunk_voxel_count::<CBITS>() {
                        d.push(p.get_tile_material(t));
                    }
                    *self = TileMaterial::Direct(d);
                    if let TileMaterial::Direct(v) = self {
                        if let Some(mat) = m {
                            v[i] = Some(mat);
                        } else {
                            v[i] = None;
                        }
                    }
                }
            }
            TileMaterial::Direct(v) => {
                if let Some(mat) = m {
                    v[i] = Some(mat);
                } else {
                    v[i] = None;
                }
            }
        }
    }
}
